import time
from contextlib import contextmanager
from flask_sqlalchemy import SQLAlchemy as _SQLAlchemy, BaseQuery
from sqlalchemy import Column, String


class SQLAlchemy(_SQLAlchemy):
    @contextmanager
    def auto_commit(self):
        # 开启事务,自动提交和回滚
        try:
            yield
            self.session.commit()
        except Exception as e:
            self.session.rollback()
            raise e


class Query(BaseQuery):
    def filter_by(self, **kwargs):
        # if 'status' not in kwargs.keys():
        #     kwargs['status'] = 1
        # return super(Query, self).filter_by(**kwargs)
        return super(Query, self).filter_by(**kwargs)

    def filter(self, *kwargs):
        return super(Query, self).filter(*kwargs)


db = SQLAlchemy(query_class=Query)


class Base(db.Model):
    __abstract__ = True
    create_time = Column(String(20), nullable=False)

    def set_attr(self, attrs_dict):
        # 自动赋值
        for key, value in attrs_dict.items():
            if hasattr(self, key) and key != 'id':
                setattr(self, key, value)

    def __init__(self):
        self.create_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())

    @staticmethod
    def model_to_dict(result):
        from collections import Iterable
        try:
            if isinstance(result, Iterable):
                tmp = [dict(zip(res.__dict__.keys(), res.__dict__.values())) for res in result]
                for t in tmp:
                    t.pop('_sa_instance_state')
            else:
                tmp = dict(zip(result.__dict__.keys(), result.__dict__.values()))
                tmp.pop('_sa_instance_state')
            return tmp
        except BaseException as e:
            raise TypeError('type error of result')

